# -*- coding: utf-8 -*-
from tkinter import Frame, Tk, BOTH, Entry, ttk,Label
from tkinter import scrolledtext
from tkinter.ttk import LabelFrame
from threading import Thread
import windnd
import os
import zipfile
import comtypes.client
import sqlite3
from datetime import datetime
import win32com.client as win32
import PyPDF2  # 用于加密 PDF 文件


class mygui():
    def __init__(self, root1):
        self.window_name = root1
        self.init_gui()
        # 拖入文件
        windnd.hook_dropfiles(self.window_name, func=self.dragged1)

    def init_gui(self):
        # 窗口名基本设置，主要调整尺寸
        self.window_name.title("++常用文档++批量加密工具   By:风二西 2025.02.23")
        self.window_name.geometry('500x200+400+50')
        self.window_name.attributes('-topmost', True)

        # 创建 Notebook 控件
        self.tab_control = ttk.Notebook(self.window_name)
        self.tab_control.pack(fill=BOTH, expand=True)

        # 输出主界面分页
        self.create_main_page()
        # 加密文件信息分页
        self.create_encrypted_info_page()

        # 绑定分页切换事件
        self.tab_control.bind("<<NotebookTabChanged>>", self.on_tab_changed)

        support_label = Label(self.window_name, text="支持PDF,DOCX,XLSX等文件格式",
                              anchor="e", foreground="blue")
        support_label.pack(side="bottom", pady=1, fill="x")
    def create_main_page(self):
        self.c_1 = Frame(self.window_name)
        self.tab_control.add(self.c_1, text='加密')

        # 第0行 设置密码
        self.r_1 = LabelFrame(self.c_1, text="设置密码")
        self.r_1.pack(fill=BOTH, pady=1, padx=1)
        self.passwd1 = Entry(self.r_1, bg="Ivory")
        self.passwd1.pack(fill=BOTH, pady=1, padx=1)

        # 第三行
        self.r_3 = LabelFrame(self.c_1, text="信息输出")
        self.r_3.pack(fill=BOTH, expand=True)
        self.output_text = scrolledtext.ScrolledText(self.r_3, bg="Ivory", font=('楷体', 10), height=5)
        self.output_text.pack(fill=BOTH, expand=True, padx=2)
        self.output_text.tag_add('tag1', "end")
        self.output_text.tag_config('tag1', foreground='red')

    def create_encrypted_info_page(self):
        self.c_2 = Frame(self.window_name)
        self.tab_control.add(self.c_2, text='数据库')

    # 窗口函数
    def dragged1(self, files):
        data = files
        self.t1 = Thread(target=self.dragged_files, args=(data,))
        self.t1.setDaemon(True)
        self.t1.start()

    # 拖入文件函数
    def dragged_files(self, files):
        # 在当前线程中创建数据库连接和游标
        conn = sqlite3.connect('encryption.db')
        cursor = conn.cursor()
        # 创建表（如果不存在），增加 encryption_time 字段
        cursor.execute('''CREATE TABLE IF NOT EXISTS encrypted_files
                               (filename TEXT, password TEXT, encryption_time TEXT)''')

        file_list = []
        for i in files:
            if os.path.isdir(i.decode("gbk")):
                for h, j, k in os.walk(i.decode("gbk")):
                    for dd in k:
                        if dd.endswith(('.docx', '.pdf', '.xlsx')):
                            file_list.append(os.path.join(h, dd))
            if os.path.isfile(i.decode("gbk")):
                if i.decode("gbk").endswith(('.docx', '.pdf', '.xlsx')):
                    file_list.append(i.decode("gbk"))

        for file_path in file_list:
            try:
                password = self.passwd1.get()
                file_ext = os.path.splitext(file_path)[1].lower()
                if file_ext == '.docx':
                    if self.is_docx_protected(file_path):
                        self.set_docx_password(file_path, password)
                        self.output_text.insert("end", f"{file_path} 完成加密\n")
                    else:
                        self.output_text.insert("end", f"{file_path} 存在密码\n", "tag1")
                        continue
                elif file_ext == '.pdf':
                    if not self.is_pdf_protected(file_path):
                        if self.set_pdf_password(file_path, password):
                            self.output_text.insert("end", f"{file_path} 完成加密\n")
                        else:
                            continue
                    else:
                        self.output_text.insert("end", f"{file_path} 存在密码\n", "tag1")
                        continue
                elif file_ext == '.xlsx':
                    if self.is_xlsx_protected(file_path):
                        self.set_xlsx_password(file_path, password)
                        self.output_text.insert("end", f"{file_path} 完成加密\n")
                    else:
                        self.output_text.insert("end", f"{file_path} 存在密码\n", "tag1")
                        continue

                # 提取文件名
                file_name = os.path.basename(file_path)
                # 获取当前时间
                encryption_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                # 将文件名、密码和加密时间存入数据库
                cursor.execute("INSERT INTO encrypted_files VALUES (?,?,?)", (file_name, password, encryption_time))
                conn.commit()
            except Exception as e:
                self.output_text.insert("end", f"{file_path} 发生错误: {e}\n")
        # 关闭连接
        conn.close()

    @staticmethod
    def thread123(func, *args):
        '''将函数打包进线程'''
        # 创建
        t = Thread(target=func, args=args)
        # 守护 !!!
        t.setDaemon(True)
        # 启动
        t.start()
        # 阻塞--卡死界面！
        # t.join()

    def on_tab_changed(self, event):
        # 获取当前选中的分页索引
        current_tab_index = self.tab_control.index(self.tab_control.select())
        # 检查是否切换到了“加密文件信息”分页
        if self.tab_control.tab(current_tab_index, "text") == '加密文件信息':
            self.load_encrypted_files()

    def load_encrypted_files(self):
        # 在当前线程中创建数据库连接和游标
        conn = sqlite3.connect('encryption.db')
        cursor = conn.cursor()

        # 清空加密文件信息分页的内容
        for widget in self.c_2.winfo_children():
            widget.destroy()

        # 查询数据库中的数据，包含加密时间
        cursor.execute("SELECT * FROM encrypted_files")
        rows = cursor.fetchall()

        # 创建 Treeview 表格
        columns = ('文件名', '密码', '加密时间')
        tree = ttk.Treeview(self.c_2, columns=columns, show='headings')
        tree.pack(fill=BOTH, expand=True, padx=10, pady=10)

        # 设置列标题
        for col in columns:
            tree.heading(col, text=col)
            tree.column(col, width=150)

        # 插入数据
        for row in rows:
            tree.insert('', 'end', values=row)

        # 关闭连接
        conn.close()

    # 检测 word 是否已经加密
    def is_docx_protected(self, file_path):
        try:
            with zipfile.ZipFile(file_path, 'r') as zip_ref:
                if 'encryption.xml' in zip_ref.namelist():
                    return 0
        except Exception as e:
            return 0
        return 1

    def set_docx_password(self, file_path, password):
        word = comtypes.client.CreateObject('Word.Application')
        word.Visible = False
        word.DisplayAlerts = False
        doc = word.Documents.Open(file_path)
        doc.Password = password
        doc.Save()
        doc.Close()
        word.Quit()

    # 检测 PDF 是否已经加密
    def is_pdf_protected(self, file_path):
        try:
            with open(file_path, 'rb') as infile:
                reader = PyPDF2.PdfReader(infile)
                return reader.is_encrypted
        except Exception as e:
            return 0

    def set_pdf_password(self, input_pdf, password):
        """
        为指定的 PDF 文件设置打开密码。
        :param input_pdf: 输入的 PDF 文件路径
        :param password: 密码字符串
        """
        try:
            # 打开输入的 PDF 文件
            with open(input_pdf, 'rb') as infile:
                reader = PyPDF2.PdfReader(infile)

                # 创建一个 PdfWriter 对象
                writer = PyPDF2.PdfWriter()

                # 将所有页面添加到 PdfWriter 对象
                for page_num in range(len(reader.pages)):
                    writer.add_page(reader.pages[page_num])

                # 设置打开密码
                writer.encrypt(user_password=password, owner_password=None, use_128bit=True)

                # 写入带密码的 PDF 文件
                with open(input_pdf, 'wb') as outfile:
                    writer.write(outfile)

            return 1
        except Exception as e:
            self.output_text.insert("end", f"{input_pdf} PDF 加密失败: {e}\n", "tag1")
            return 0

    # 检测 XLSX 是否已经加密
    def is_xlsx_protected(self, file_path):
        try:
            excel = win32.gencache.EnsureDispatch('Excel.Application')
            excel.DisplayAlerts = False
            wb = excel.Workbooks.Open(file_path)
            wb.Close()
            excel.Quit()
            return 1
        except Exception:
            return 0

    def set_xlsx_password(self, file_path, password):
        try:
            excel = win32.gencache.EnsureDispatch('Excel.Application')
            excel.DisplayAlerts = False
            wb = excel.Workbooks.Open(file_path)
            wb.Password = password
            wb.Save()
            wb.Close()
            excel.Quit()
        except Exception as e:
            self.output_text.insert("end", f"{file_path} XLSX 加密失败: {e}\n", "tag1")


def gui_start():
    root = Tk()  # 实例化出一个父窗口
    scanqr = mygui(root)
    root.mainloop()  # 父窗口进入事件循环，可以理解为保持窗口运行，否则界面不展示


gui_start()